# coding=utf-8
import pandas as pd
import matplotlib.pyplot as plt
import time
import numpy as np
import os
import random
from math import sin
from math import cos
from math import sqrt
from time import sleep


class Swarm:
    def __init__(self, num=4):
        self.path = os.path.abspath('..')
        self.agent_num = num
        self.size = self.agent_num
        self.K = np.array([0.213, 0.0952])
        self.H = np.array(
            [[0, 0, 1, 1],
             [10, 0, 1, 1],
             [10, 10, 1, 1],
             [0, 10, 1, 1]]
        )
        self.agent_safe_r = 0.5
        self.safe_move_area_x = [-20, 20]
        self.safe_move_area_y = [-20, 20]
        self.v_max = 10
        self.step = 0.1  # 时间
        self.min_v = 0.001  # 移动最小速度
        self.global_position = np.zeros((self.agent_num, 2))
        self.global_vel = np.zeros((self.agent_num, 2))
        self.PI = 3.1415926

        self.adj_matrix = np.ones((self.size, self.size))
        for i in range(self.size):
            self.adj_matrix[i][i] = 0

        pose_init_ok = 0
        while not pose_init_ok:
            for i in range(self.agent_num):
                self.global_position[i][0] = random.uniform(self.safe_move_area_x[0] / 2, self.safe_move_area_x[1] / 2)
                self.global_position[i][1] = random.uniform(self.safe_move_area_x[0] / 2, self.safe_move_area_x[1] / 2)
            for i in range(self.agent_num):
                dist = self.get_distance(self.global_position, i)
                for distance in dist:
                    if distance < 2 * self.agent_safe_r:
                        pose_init_ok = 0
                        break
            pose_init_ok = 1
        # 生成画布
        plt.figure(figsize=(8, 6), dpi=80)
        # 打开交互模式
        plt.ion()

    def get_distance(self, target_position, index):
        res = np.zeros(self.size)
        for i in range(self.size):
            res[i] = res[i] + sqrt((target_position[i][0] - target_position[index][0]) *
                                   (target_position[i][0] - target_position[index][0]) +
                                   (target_position[i][1] - target_position[index][1]) *
                                   (target_position[i][1] - target_position[index][1]))
        return res

    # 指定变化index在addedindex 对应index之后比如
    def added_agent_swarm(self, added_num, added_position=None, added_agent_adj=None):
        if added_position is None:
            added_position = []
        if added_agent_adj is None:
            added_agent_adj = []
        self.agent_num = self.agent_num + added_num
        self.size = self.agent_num
        new_global_position = self.global_position.copy()
        if len(added_position) == added_num:
            for i in range(added_num):
                new_global_position = np.insert(new_global_position, self.agent_num - added_num + i, added_position[i],
                                                0)
        else:
            added_position = np.zeros((added_num, 2))
            init_ok = 0
            while not init_ok:
                for i in range(added_num):
                    added_position[i][0] = random.uniform(self.safe_move_area_x[0], self.safe_move_area_x[1])
                    added_position[i][1] = random.uniform(self.safe_move_area_x[0], self.safe_move_area_x[1])
                for i in range(added_num):
                    new_global_position = np.insert(new_global_position, self.agent_num - added_num + i,
                                                    added_position[i],
                                                    0)
                for i in range(self.agent_num):
                    dist = self.get_distance(new_global_position, i)
                    for distance in dist:
                        if distance < 2 * self.agent_safe_r:
                            init_ok = 0
                            break
                init_ok = 1
        new_adj_matrix = np.zeros((self.size, self.size))
        if len(added_agent_adj) != added_num or len(added_agent_adj[0]) != size:
            added_agent_adj = np.ones((added_num, self.size))
            for i in range(added_num):
                added_agent_adj[i][i + self.agent_num - added_num] = 0

        for i in range(len(self.adj_matrix)):
            for j in range(len(self.adj_matrix)):
                new_adj_matrix[i][j] = self.adj_matrix[i][j]

        for i in range(added_num):
            new_adj_matrix[i + self.agent_num - added_num, :] = added_agent_adj[i]
            new_adj_matrix[:, i + self.agent_num - added_num] = added_agent_adj[i]

        self.global_position = new_global_position.copy()
        self.adj_matrix = new_adj_matrix.copy()

    def remove_swarm_index(self, index):
        self.agent_num = self.agent_num - 1
        self.size = self.size - 1
        self.adj_matrix = np.delete(self.adj_matrix, index, 0)
        new_adj_matrix = np.delete(self.adj_matrix, index, 1)
        self.adj_matrix = new_adj_matrix
        new_global_position = np.delete(self.global_position, index, 0)
        self.global_position = new_global_position

    def get_degree_matrix(self, adj_m):
        res = np.zeros((self.size, self.size))
        for i in range(self.size):
            degree = 0
            for j in range(self.size):
                degree = degree + adj_m[i][j]
            res[i][i] = degree
        return res

    def get_L_matrix(self, adj_m):
        return self.get_degree_matrix(adj_m) - adj_m

    def plot_adj_matrix(self, adj_m, plot_r=1):
        points = np.zeros((self.size, 2))
        for i in range(self.size):
            points[i][0] = plot_r * sin(2 * self.PI / self.size / 2 * (2 * i + 1))
            points[i][1] = plot_r * cos(2 * self.PI / self.size / 2 * (2 * i + 1))
        plt.plot(points[:, 0], points[:, 1], 'ro')
        for i in range(self.size):
            for j in range(self.size):
                if j >= i:
                    if adj_m[i][j] > 0:
                        plt.plot([points[i][0], points[j][0]], [points[i][1], points[j][1]], 'r')
                else:
                    continue
        for i in range(self.size):
            plt.text(points[i][0], points[i][1], i + 1, fontsize=20, color='black')
        plt.axis([-plot_r * 4.0 / 3, plot_r * 4.0 / 3, -plot_r * 4.0 / 3, plot_r * 4.0 / 3])
        plt.title("Agents Adjacency Graph")
        plt.ylabel("")
        plt.xlabel("")
        # plt.grid(True)
        plt.show()

    def control_by_index(self, index):
        x = self.global_position[index][0]
        y = self.global_position[index][1]
        if x < self.safe_move_area_x[0]:
            print "x out of safe"
            return [self.v_max, 0]
        if x > self.safe_move_area_x[1]:
            print "x out of safe"
            return [-self.v_max, 0]
        if y < self.safe_move_area_y[0]:
            print "y out of safe"
            return [0, self.v_max]
        if y > self.safe_move_area_y[1]:
            print "y out of safe"
            return [0, -self.v_max]
        vx = 0
        vy = 0
        cur_distance = self.get_distance(self.global_position, index)
        not_safe = 1
        while not_safe:
            not_safe = 0
            target_position = self.global_position.copy()
            target_vel = self.global_vel.copy()
            for i in range(self.size):
                dx = (target_position[i][0] - self.H[i][0]) - (target_position[index][0] - self.H[index][0])
                dy = (target_position[i][1] - self.H[i][1]) - (target_position[index][1] - self.H[index][1])
                dvx = (target_vel[i][0] - self.H[i][2]) - (target_vel[index][0] - self.H[index][2])
                dvy = (target_vel[i][1] - self.H[i][3]) - (target_vel[index][1] - self.H[index][3])
                if self.adj_matrix[index][i] == 0:
                    continue
                if 2 * self.agent_safe_r > cur_distance[i] > 0:
                    # TODO: 可以考虑速度进去，人工势场法
                    vx = vx + 5 * -dx * pow((2 * self.agent_safe_r - sqrt(dx * dx + dy * dy)), 2) / sqrt(
                        dx * dx + dy * dy)
                    vy = vy + 5 * -dy * pow((2 * self.agent_safe_r - sqrt(dx * dx + dy * dy)), 2) / sqrt(
                        dx * dx + dy * dy)
                    continue
                vx = vx + self.K[0] * dx + self.K[1] * dvx
                vy = vy + self.K[0] * dy + self.K[1] * dvy
            if sqrt(vx * vx + vy * vy) > self.v_max:
                _vx = self.v_max * vx / sqrt(vx * vx + vy * vy)
                _vy = self.v_max * vy / sqrt(vx * vx + vy * vy)
                vx = _vx
                vy = _vy

            target_position[index][0] = target_position[index][0] + vx * self.step
            target_position[index][1] = target_position[index][1] + vy * self.step
            dist = self.get_distance(target_position, index)
            for i in range(self.size):
                if dist[i] > self.agent_safe_r * 2:
                    not_safe = 0
            if not_safe:
                vx = vx / 4
                vy = vy / 4
                if abs(vx) < self.min_v:
                    vx = 0
                if abs(vy) < self.min_v:
                    vy = 0
                if abs(vx) < self.min_v and abs(vy) < self.min_v:
                    return [0, 0]
            else:
                break
        self.global_vel[index][0] = vx
        self.global_vel[index][1] = vy
        print index, vx, vy
        return [vx, vy]
        pass

    def plot_swarm(self):
        # plt.ion()：打开交互模式
        # plt.ioff()：关闭交互模式
        # plt.clf()：清除当前的Figure对象
        # plt.cla()：清除当前的Axes对象
        # plt.pause()：暂停功能
        plt.cla()

        for i in range(self.size):
            plt.plot(self.global_position[:, 0], self.global_position[:, 1], 'bo')
            plt.text(self.global_position[i][0], self.global_position[i][1], i + 1, fontsize=16, color='black')
            for j in range(self.size):
                if j >= i:
                    if self.adj_matrix[i][j] == 0:
                        continue
                    plt.plot([self.global_position[i, 0], self.global_position[j, 0]],
                             [self.global_position[i, 1], self.global_position[j, 1]],
                             'r')

        plt.axis(
            [self.safe_move_area_x[0], self.safe_move_area_x[1], self.safe_move_area_y[0], self.safe_move_area_y[1]])
        plt.title("Swarm nodes")
        plt.ylabel("")
        plt.xlabel("")
        # plt.grid(True)
        plt.show()
        plt.pause(self.step)

    # 简单的一致性控制，不考虑进出节点
    def control_simple(self):
        while 1:
            v = np.zeros((self.size, 2))
            for i in range(4):
                v[i] = self.control_by_index(i)
            # print v
            for i in range(4):
                self.global_position[i][0] = self.global_position[i][0] + v[i][0] * self.step
                self.global_position[i][1] = self.global_position[i][1] + v[i][1] * self.step
            sleep(self.step)
            self.plot_swarm()

    # 简单的一致性控制，考虑移除节点
    def control_remove(self):
        count = 0
        remove_flag = 0
        random_count = random.randint(10, 100)
        while 1:
            count = count + 1
            v = np.zeros((self.size, 2))
            for i in range(self.size):
                v[i] = self.control_by_index(i)
            # print v
            for i in range(self.size):
                self.global_position[i][0] = self.global_position[i][0] + v[i][0] * self.step
                self.global_position[i][1] = self.global_position[i][1] + v[i][1] * self.step
            sleep(self.step)
            self.plot_swarm()
            if count == random_count and not remove_flag:
                remove_flag = 1
                self.remove_swarm_index(random.randint(0, self.size - 1))

    # 简单的一致性控制，考虑添加节点
    def control_add(self, num):
        count = 0
        add_flag = 0
        random_count = random.randint(10, 100)
        while 1:
            count = count + 1
            v = np.zeros((self.size, 2))
            for i in range(self.size):
                v[i] = self.control_by_index(i)
            # print v
            for i in range(self.size):
                self.global_position[i][0] = self.global_position[i][0] + v[i][0] * self.step
                self.global_position[i][1] = self.global_position[i][1] + v[i][1] * self.step
            sleep(self.step)
            self.plot_swarm()
            if count == random_count and not add_flag:
                add_flag = 1
                self.added_agent_swarm(num)

    # 简单的一致性控制，考虑添加与移除节点
    def control_add_remove(self, add_num, remove_index):
        count = 0
        add_flag = 0
        remove_flg = 0
        random_count1 = random.randint(10, 100)
        random_count2 = random.randint(10, 100)
        while 1:
            count = count + 1
            v = np.zeros((self.size, 2))
            for i in range(self.size):
                v[i] = self.control_by_index(i)
            # print v
            for i in range(self.size):
                self.global_position[i][0] = self.global_position[i][0] + v[i][0] * self.step
                self.global_position[i][1] = self.global_position[i][1] + v[i][1] * self.step
            sleep(self.step)
            self.plot_swarm()
            if count == random_count1 and not add_flag:
                add_flag = 1
                self.added_agent_swarm(add_num)
            if count == random_count2 and not remove_flg:
                remove_flg = 1
                self.remove_swarm_index(remove_index)


if __name__ == "__main__":
    simulation_swarm = Swarm()
    simulation_swarm.control_simple()
    # simulation_swarm.control_add(1)
    # simulation_swarm.control_remove()
    # print simulation_swarm.global_position
    # print simulation_swarm.adj_matrix
    # simulation_swarm.control_all()
    # simulation_swarm.added_agent_swarm(3)
    # simulation_swarm.remove_swarm_index(2)
    # print simulation_swarm.global_position
    # print simulation_swarm.adj_matrix
    # simulation_swarm.control_add_remove(2, 3)
    pass
    # plot_adj_matrix(adj_matrix)
